AIP-104: Task Iteration - #62922

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

Claude Code (Fable 5.1).

Generated-by: Claude Code (Fable 5.1) following the guidelines

Description

This PR is the initial implementation of Iterable Tasks (IT), as discussed in the devlist and building upon the foundations of AIP-104. (Originally prototyped as "Dynamic Task Iteration"; renamed to Iterable Tasks following review feedback to avoid confusion with Dynamic Task Mapping.)

For further context on the use cases and performance benefits of IT, see this Medium Article and the new dynamic-task-mapping-vs-iteration.rst doc added in this PR, which compares IT with Dynamic Task Mapping (DTM) and Dynamic Task Batching in depth.

The XCom Database Constraint Challenge

While porting our internal "monkey-patched" version of IT (used since Airflow 2.x) to the core, I've identified a significant technical hurdle regarding XCom handling.

Around Airflow 2.10/2.11, a change was introduced to the database constraints for the XCom table. Specifically:

• Current State: The DB prevents creating indexed XComs (map_index >= 0) unless a corresponding mapped TaskInstance exists in the task_instance table.
• The Conflict: IT is designed to process multiple indexed XComs within a single Task Instance. Because there is no 1-to-1 mapping of a sub-task index to a physical TI row, the DB constraint blocks the insertion of these results.

The drawback is that XComs wouldn't automatically be removed from the database when a TaskInstance is deleted, which is the purpose of that constraint. So appending the index to the XCom key would be a good enough solution for IT, but not for DTM.

Current implementation in this PR

• XComIterable (airflow.sdk.bases.xcom) appends the sub-task index directly to the XCom key (return_value_<index>) to bypass the constraint, and exposes the results as a lazy Sequence (__len__/__getitem__/__iter__) so a downstream task can consume them the same way it would consume an .expand() result.
• Progress and crash recovery no longer rely on XCom alone. IterableOperator tracks per-sub-task progress in the AIP-103 Task State Store rather than XCom, and participates in Airflow's standard retry mechanism: if the IterableOperator's task instance is retried (or manually cleared before it finishes), already-succeeded sub-tasks are skipped and only pending/failed ones re-run. XCom is only (re)written per index once a sub-task succeeds; once every index has succeeded, a completion marker is written so that a subsequent manual clear reruns all indices from scratch rather than replaying stale results. The checkpoints themselves are kept for the store's retention ([state_store] default_retention_days), so a manual clear of a task that exhausted its retries resumes from them.
• Outlet asset/inlet events and on_kill propagation are handled per sub-task: a checkpointed sub-task replays its recorded outlet events on the following attempt instead of losing them, and killing the IterableOperator's task instance propagates to any sub-tasks still in flight.

I believe the cleanest long-term path is still to add a dedicated route in the Execution API that retrieves multiple XComs for a single TaskInstance by a list of keys in one round trip, so XComIterable.__getitem__/slicing don't need one request per element. I have a PR open to address this, intentionally split out of this PR.

This was also discussed in the devcall, see 2026-06-04 Dev Call Minutes.

AIP-104 itself was discussed again in the latest devcall, where the concerns raised there have also been addressed: 2026-09-10 Dev call Minutes.

Examples

The examples below assume an HTTP connection named pokeapi pointing to https://pokeapi.co.

Task Iteration

This example fetches a list of Pokémon from the PokéAPI and then uses Iterable Tasks (IT) to retrieve the details of each Pokémon. A single task instance processes all Pokémon URLs.

from airflow.sdk import dag, task
from airflow.providers.http.hooks.http import HttpHook, HttpAsyncHook

from pendulum import datetime

@dag(
    start_date=datetime(2025, 1, 1),
    schedule=None,
    catchup=False,
)
def pokemon_iteration():
    @task
    def list_pokemon() -> list[str]:
        response = HttpHook(
            http_conn_id="pokeapi",
            method="GET",
        ).run(
            endpoint="api/v2/pokemon?limit=100",
        )

        return [
            pokemon["url"].replace("https://pokeapi.co/", "")
            for pokemon in response.json()["results"]
        ]

    @task(
        retries=3,
        task_concurrency=2,
        show_return_value_in_logs=False,
    )
    async def get_pokemon(url: str):
        async with HttpAsyncHook(
            http_conn_id="pokeapi",
            method="GET",
        ).session() as session:
            response = await session.run(endpoint=url)
            return await response.json()

    get_pokemon.iterate(
        url=list_pokemon(),
    )

pokemon_iteration()

Comparison

This demonstrates how Task Iteration can significantly reduce TaskInstance creation overhead. Task Spreading (running one iteration over exactly N TaskInstances with .spread(across=N).iterate()) is split out into #73688.

Notable design points addressed since the initial draft

• .iterate()'s dict-argument semantics now match .expand(): passing a dict value forwards (key, value) pairs to each sub-task instead of bare keys.
• IterableOperator.operator_name forwards to the wrapped operator (including @task-decorated callables with a custom_operator_name), so sub-tasks report the correct type in the UI/API instead of always showing MappedOperator/IterableOperator.
• XComIterable.flatten() moved out to Add XComIterable.flatten() to read an iterated task's pages as one sequence #73807, stacked on this PR, so this PR stays about running a task over its input.
• .iterate() refuses a literal that is no collection (a scalar, a string, None) at parse time on both the classic and the @task path, as .expand() does, and the same value from an upstream at run time.
• execution_timeout caps the whole iteration; no limit applies per indexed task, sync or async (one of the same value, started later, could never fire first). An AirflowTaskTimeout a sync indexed task raises itself (a hook that gave up waiting) is that task's own failure, as it is the mapped task instance's under .expand().
• on_kill() takes no lock, so the runner's SIGTERM handler cannot hang on the loop thread's own register/unregister; the copy prepare_for_execution makes gets an iteration state of its own, so a kill in one attempt does not stop the next under dag.test(); an asend cancelled while waiting for the comms thread lock no longer leaves it taken (CommsDecoder._acquire_thread_lock).
• Operators that find their remote work by the task instance's identity (KubernetesPodOperator reattaching, EcsRunTaskOperator(reattach=True)) are listed under "When not to use IT", with {{ ti.index }} as the label that tells items apart.
• on_kill() propagates to in-flight sub-tasks, stops the iteration from starting any other (one pulled before the kill but not yet started is counted apart and left for the next attempt), and outlet/asset events recorded by a sub-task that already succeeded are replayed from its checkpoint on a later retry instead of being lost.
• Deferred operators, reschedule-mode sensors, TriggerDagRunOperator, and ShortCircuitOperator-style downstream skipping are explicitly rejected inside IterableOperator with actionable errors rather than being silently mishandled — see the class docstring for the full list of current limitations.
• multiple_outputs is ignored for iterated tasks, explicitly at the IterableOperator level. A @task with a Mapping return annotation infers multiple_outputs=True, but the value the runner pushes for an iterated task is the XComIterable aggregate rather than a dict, so honouring the flag made the runner reject the result after every sub-task had already succeeded. Each sub-task's return value is pushed whole as return_value_<index>; keys are not fanned out into separate XComs the way .expand() does. Documented on the class and in the Task SDK docs, and pinned by a runner-level regression test for a dict-returning task under .iterate().

Per-iteration keys: XComs and task state

Every iteration of an iterated task runs under the same task instance (same dag id, task id, run id and map index). Anything an iteration writes into a per-task-instance store therefore competes with its siblings for the same key, and with the async executor the winner is whichever iteration finishes last. Two stores are affected, and both now apply the same rule: a key written from inside an iteration carries that iteration's index.

• XComs. IndexedTaskInstance.xcom_push/axcom_push suffix the key with _<index>. That is what makes return_value_<index> and XComIterable work, and it applies to any key an operator pushes from execute, including the keys of a multiple_outputs dict. Pulls are not suffixed: ti.xcom_pull(task_ids="upstream") reaches the upstream's XCom untouched.
• Task state store (the AIP-103 store, context["task_state_store"]). This was a gap: an iteration that stored a watermark or a cursor with task_state_store.set("last_offset", ...) shared that key with every sibling. IndexedTaskStateStoreAccessor closes it: IndexedTaskInstance.task_state_store is the parent's accessor seen through the index, suffixing keys on get/set/delete and their async twins, and the sub-task's context carries the same object, so an operator does not need to know it is being iterated. clear() is refused inside an iteration, since it would wipe the siblings' state and the operator's own checkpoints; an iteration deletes its own keys instead.
• The operator's checkpoints are separate. IterableOperator records per-index progress in the parent's store under _iterable_<index> and _iterable_completed, written through the parent's accessor, so they are never double-suffixed and never collide with user keys.

IndexedTaskRunner (formerly TaskExecutor, renamed because it read like one of Airflow's executors) builds the context an iteration runs against: a copy of the parent's context with the iteration's own task instance, its indexed state store view and its own outlet events. The operator binds it from the with statement, runs the operator inside that block, and records the outcome (checkpoint, XCom push, outlet-event merge) after it, so on_kill and the failure callbacks apply to the operator's execution only.

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.